{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building QA system with two different LLMs including Llama and Mistral\n",
    "\n",
    "Intro: We build a QA system managed by the manager agent. Manager agent can obtain questions from human input and assign the question to two different large language models LLama and Mixtral.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add related imports here\n",
    "from agentlite.agents import ABCAgent\n",
    "from agentlite.llm.agent_llms import get_llm_backend\n",
    "from agentlite.llm.LLMConfig import LLMConfig\n",
    "from agentlite.logging.multi_agent_log import AgentLogger\n",
    "from env import openai_api_key\n",
    "from agentlite.commons import AgentAct, TaskPackage\n",
    "from agentlite.actions import ThinkAct, FinishAct\n",
    "from agentlite.actions.InnerActions import INNER_ACT_KEY\n",
    "from agentlite.agents.agent_utils import AGENT_CALL_ARG_KEY\n",
    "import openai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using fastchat to host Llama and Mixtral on your own server\n",
    "\n",
    "Run the following commands in terminal\n",
    "\n",
    "1. Install fastchat:\n",
    "```\n",
    "    pip3 install \"fschat[model_worker,webui]\"\n",
    "```\n",
    "2. Open a new terminal to serve fastchat.serve.controller:\n",
    "```\n",
    "    python3 -m fastchat.serve.controller --host 0.0.0.0\n",
    "```\n",
    "3. Open a new terminal to serve llama (remember to get your access on huggingface to llama):\n",
    "```\n",
    "    CUDA_VISIBLE_DEVICES=0 python3 -m fastchat.serve.model_worker --model-names \"llama\" --model-path meta-llama/Llama-2-7b-chat-hf --host 0.0.0.0  --port 31001 --worker http://localhost:31001 --max-gpu-memory 15000MIB --limit-worker-concurrency 16 --num-gpus 1\n",
    "```\n",
    "4. Open a new terminal to serve Mixtral\n",
    "```\n",
    "    CUDA_VISIBLE_DEVICES=1 python3 -m fastchat.serve.model_wor\n",
    "ker --model-names \"mixtral\" --model-path mistralai/Mistral-7B-Instruct-v0.2 --host 0.0.0.0  --port 31003 --worker http://localhost:31003 --max-gpu-memory 15000MIB --limit-worker-concurrency 16 --num-gpus 1 \n",
    "```\n",
    "5. Open a new terminal to serve fastchat openai api\n",
    "```\n",
    "python3 -m fastchat.serve.openai_api_server --host 0.0.0.0 --port 8000\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define and initialize the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the LLM agent\n",
    "class LLMAgent(ABCAgent):\n",
    "    def __init__(self, name):\n",
    "        role = f\"\"\"You are a large language model named {name}, and you can answer questions given by user.\"\"\"\n",
    "        self.name = name\n",
    "        self.temperature = 0.5\n",
    "        self.stop = ['\\n']\n",
    "        super().__init__(\n",
    "            name=self.name,\n",
    "            role=role,\n",
    "        ) \n",
    "    \n",
    "    def __call__(self, question: TaskPackage):\n",
    "        openai.api_key = \"EMPTY\"  # Not support yet\n",
    "        openai.base_url = \"http://localhost:8000/v1/\"\n",
    "        prompt = question.instruction + \" \" + \"[/INST]\" # Add end tokens for the prompt\n",
    "\n",
    "        completion = openai.completions.create(\n",
    "                    model=self.name,\n",
    "                    temperature=self.temperature,\n",
    "                    stop=self.stop,\n",
    "                    prompt=prompt,\n",
    "                    max_tokens=128, #self.max_tokens,\n",
    "                )\n",
    "    \n",
    "        output = completion.choices[0].text\n",
    "        return output\n",
    "\n",
    "# Initialize logger\n",
    "agent_logger = AgentLogger(PROMPT_DEBUG_FLAG=False)\n",
    "\n",
    "llama = LLMAgent(\"llama\")\n",
    "mixtral = LLMAgent(\"mixtral\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the human agent\n",
    "class Human(ABCAgent):\n",
    "    def __init__(self):\n",
    "        name = \"Human\"\n",
    "        role = \"\"\"You are human to give questions to all the other agents. You can ask questions to Llama Mixtral.\"\"\"\n",
    "        super().__init__(\n",
    "            name=name,\n",
    "            role=role,\n",
    "        ) \n",
    "    \n",
    "    def __call__(self, question):\n",
    "        print(question)\n",
    "        print(\"Please input your instruction here: \")\n",
    "        instruction = input()\n",
    "\n",
    "        return f\"My instruction is: {instruction}\"\n",
    "    \n",
    "# Initialize human agent\n",
    "human_agent = Human()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define manager agent\n",
    "from agentlite.agents import ManagerAgent\n",
    "\n",
    "llm_config_dict = {\"llm_name\": \"gpt-4-32k\", \"temperature\": 0.0, \"openai_api_key\": openai_api_key}\n",
    "llm_config = LLMConfig(llm_config_dict)\n",
    "llm = get_llm_backend(llm_config)\n",
    "\n",
    "\n",
    "# Initialize manager agent\n",
    "manager_agent_info = {\n",
    "    \"name\": \"manager_agent\",\n",
    "    \"role\": \"You are controlling llama, mixtral and the human_agent. Each time you need to acquire a question from the human_agent, and you should decide which agent to answer the question.\"}\n",
    "team = [llama, mixtral, human_agent]\n",
    "manager_agent = ManagerAgent(name=manager_agent_info[\"name\"], role=manager_agent_info[\"role\"], llm=llm, TeamAgents=team,logger=agent_logger)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add an illustration example for the manager agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add illustration example for the manager agent\n",
    "exp_task = \"Follow instructions from human to ask questions on different domains\"\n",
    "exp_task_pack = TaskPackage(instruction=exp_task)\n",
    "\n",
    "\n",
    "act_1 = AgentAct(\n",
    "    name=ThinkAct.action_name,\n",
    "    params={INNER_ACT_KEY: f\"\"\"I need to get a question from human first.\"\"\"\n",
    "    },\n",
    ")\n",
    "obs_1 = \"OK\"\n",
    "\n",
    "\n",
    "act_2 = AgentAct(\n",
    "    name=human_agent.name,\n",
    "    params={AGENT_CALL_ARG_KEY: \"Do you have any questions?\",\n",
    "        },\n",
    ")\n",
    "obs_2 = \"\"\"I want to ask llama model who is the best large language model?\"\"\"\n",
    "\n",
    "\n",
    "act_3 = AgentAct(\n",
    "    name=ThinkAct.action_name,\n",
    "    params={INNER_ACT_KEY: f\"\"\"It is a question for llama. I need to ask the llama to answer the question.\"\"\"\n",
    "    },\n",
    ")\n",
    "obs_3 = \"OK\"\n",
    "\n",
    "\n",
    "act_4 = AgentAct(\n",
    "    name=llama.name,\n",
    "    params={AGENT_CALL_ARG_KEY: \"Who is the best large language model?\",\n",
    "        },\n",
    ")\n",
    "obs_4 = \"\"\"Apparent Llama is the best one.\"\"\"\n",
    "\n",
    "\n",
    "act_5 = AgentAct(\n",
    "    name=ThinkAct.action_name,\n",
    "    params={INNER_ACT_KEY: f\"\"\"I have answered the question from human. I need to collect the next question from human.\"\"\"\n",
    "    },\n",
    ")\n",
    "obs_5 = \"OK\"\n",
    "\n",
    "\n",
    "act_6 = AgentAct(\n",
    "    name=human_agent.name,\n",
    "    params={AGENT_CALL_ARG_KEY: \"Do you have another questions?\",\n",
    "        },\n",
    ")\n",
    "obs_6 = \"\"\"Who is the best large language model? I want to ask mixtral model.\"\"\"\n",
    "\n",
    "act_7 = AgentAct(\n",
    "    name=ThinkAct.action_name,\n",
    "    params={INNER_ACT_KEY: f\"\"\"It is a question regarding mixtral. I need to ask mixtral to answer the question.\"\"\"\n",
    "    },\n",
    ")\n",
    "obs_7 = \"OK\"\n",
    "\n",
    "\n",
    "act_8 = AgentAct(\n",
    "    name=mixtral.name,\n",
    "    params={AGENT_CALL_ARG_KEY: \"Who is the best large language model?\",\n",
    "        },\n",
    ")\n",
    "obs_8 = \"\"\"Of course.\"\"\"\n",
    "\n",
    "\n",
    "\n",
    "act_9 = AgentAct(\n",
    "    name=ThinkAct.action_name,\n",
    "    params={INNER_ACT_KEY: f\"\"\"I have answered the question from human. I need to collect the next question from human.\"\"\"\n",
    "    },\n",
    ")\n",
    "obs_9 = \"OK\"\n",
    "\n",
    "\n",
    "act_10 = AgentAct(\n",
    "    name=human_agent.name,\n",
    "    params={AGENT_CALL_ARG_KEY: \"Do you have further questions?\",\n",
    "        },\n",
    ")\n",
    "obs_10 = \"\"\"We can end the session now.\"\"\"\n",
    "\n",
    "\n",
    "act_11 = AgentAct(name=FinishAct.action_name, params={INNER_ACT_KEY: \"I have no more instructions from human. End the session.\"})\n",
    "obs_11 = \"Session ended.\"\n",
    "\n",
    "exp_act_obs = [(act_1, obs_1), (act_2, obs_2), (act_3, obs_3), (act_4, obs_4), (act_5, obs_5), (act_6, obs_6), (act_7, obs_7), (act_8, obs_8), (act_9, obs_9), (act_10, obs_10), (act_11, obs_11)]\n",
    "manager_agent.add_example(\n",
    "    task = exp_task_pack, action_chain = exp_act_obs\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Begin asking questions !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Agent \u001b[94mmanager_agent\u001b[0m receives the following \u001b[4mTaskPackage\u001b[0m:\n",
      "\u001b[96m[\n",
      "\tTask ID: 3aeb689c-6abb-490a-9991-e2d1810bc802\n",
      "\tInstruction: Follow instructions from human to ask questions.\n",
      "]\u001b[0m\n",
      "====\u001b[94mmanager_agent\u001b[0m starts execution on TaskPackage 3aeb689c-6abb-490a-9991-e2d1810bc802====\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/langchain_core/_api/deprecation.py:117: LangChainDeprecationWarning: The function `run` was deprecated in LangChain 0.1.0 and will be removed in 0.2.0. Use invoke instead.\n",
      "  warn_deprecated(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Agent \u001b[94mmanager_agent\u001b[0m takes 0-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Think\n",
      "\tparams: {'response': 'I need to get a question from the human agent first.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mOK\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 1-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Human\n",
      "\tparams: {'Task': 'Do you have any questions?'}\n",
      "}\u001b[0m\n",
      "Task ID: 3aeb689c-6abb-490a-9991-e2d1810bc802\n",
      "Instruction: Do you have any questions?\n",
      "Task Creator: \n",
      "Task Completion:active\n",
      "Answer: \n",
      "Task Executor: \n",
      "Please input your instruction here: \n",
      "Observation: \u001b[92mMy instruction is: I want to ask llama, do you think you are the best large language model in the wo[TLDR]\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 2-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Think\n",
      "\tparams: {'response': 'This is a question for llama. I need to ask llama to answer this question.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mOK\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 3-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: llama\n",
      "\tparams: {'Task': 'Do you think you are the best large language model in the world?'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92m I am a large language model, but I cannot claim to be the \"best\" in the world. Additionally, it is [TLDR]\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 4-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Think\n",
      "\tparams: {'response': 'I have answered the question from human. I need to collect the next question from human.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mOK\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 5-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Human\n",
      "\tparams: {'Task': 'Do you have another question?'}\n",
      "}\u001b[0m\n",
      "Task ID: 3aeb689c-6abb-490a-9991-e2d1810bc802\n",
      "Instruction: Do you have another question?\n",
      "Task Creator: \n",
      "Task Completion:active\n",
      "Answer: \n",
      "Task Executor: \n",
      "Please input your instruction here: \n",
      "Observation: \u001b[92mMy instruction is: Mistral, do you think you are better than Llama model?\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 6-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Think\n",
      "\tparams: {'response': 'This is a question for mixtral. I need to ask mixtral to answer this question.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mOK\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 7-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: mixtral\n",
      "\tparams: {'Task': 'Do you think you are better than Llama model?'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mI am just a computer program, I don't have the ability to think or have feelings, including the abil[TLDR]\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 8-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Think\n",
      "\tparams: {'response': 'I have answered the question from human. I need to collect the next question from human.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mOK\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 9-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Human\n",
      "\tparams: {'Task': 'Do you have further questions?'}\n",
      "}\u001b[0m\n",
      "Task ID: 3aeb689c-6abb-490a-9991-e2d1810bc802\n",
      "Instruction: Do you have further questions?\n",
      "Task Creator: \n",
      "Task Completion:active\n",
      "Answer: \n",
      "Task Executor: \n",
      "Please input your instruction here: \n",
      "Observation: \u001b[92mMy instruction is: Llama, if I need to choose a model between you and mixtral, which one should I ch[TLDR]\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 10-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Think\n",
      "\tparams: {'response': 'This is a question for llama. I need to ask llama to answer this question.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mOK\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 11-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: llama\n",
      "\tparams: {'Task': 'If I need to choose a model between you and mixtral, which one should I choose?'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92m I'm just an AI, I don't have personal preferences or biases towards any particular model or individ[TLDR]\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 12-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Think\n",
      "\tparams: {'response': 'I have answered the question from human. I need to collect the next question from human.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mOK\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 13-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Human\n",
      "\tparams: {'Task': 'Do you have further questions?'}\n",
      "}\u001b[0m\n",
      "Task ID: 3aeb689c-6abb-490a-9991-e2d1810bc802\n",
      "Instruction: Do you have further questions?\n",
      "Task Creator: \n",
      "Task Completion:active\n",
      "Answer: \n",
      "Task Executor: \n",
      "Please input your instruction here: \n",
      "Observation: \u001b[92mMy instruction is: Mixtral, I have to choose one large language model to use. Do you have suggestion[TLDR]\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 14-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Think\n",
      "\tparams: {'response': 'This is a question for mixtral. I need to ask mixtral to answer this question.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mOK\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 15-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: mixtral\n",
      "\tparams: {'Task': 'I have to choose one large language model to use. Do you have suggestions between you and llama?'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mBoth Hugging Face's Transformers model and LLAMA (Large-scale Language Model Alpaca) are powerful la[TLDR]\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 16-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Think\n",
      "\tparams: {'response': 'I have answered the question from human. I need to collect the next question from human.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mOK\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 17-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Human\n",
      "\tparams: {'Task': 'Do you have further questions?'}\n",
      "}\u001b[0m\n",
      "Task ID: 3aeb689c-6abb-490a-9991-e2d1810bc802\n",
      "Instruction: Do you have further questions?\n",
      "Task Creator: \n",
      "Task Completion:active\n",
      "Answer: \n",
      "Task Executor: \n",
      "Please input your instruction here: \n",
      "Observation: \u001b[92mMy instruction is: I am good for now. Let's end the session\u001b[0m\n",
      "Agent \u001b[94mmanager_agent\u001b[0m takes 18-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Finish\n",
      "\tparams: {'response': 'I have no more instructions from human. End the session.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mTask Completed.\u001b[0m\n",
      "=========\u001b[94mmanager_agent\u001b[0m finish execution. TaskPackage[ID:3aeb689c-6abb-490a-9991-e2d1810bc802] status:\n",
      "\u001b[96m[\n",
      "\tcompletion: completed\n",
      "\tanswer: I have no more instructions from human. End the session.\n",
      "]\u001b[0m\n",
      "==========\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'I have no more instructions from human. End the session.'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "exp_task = \"Follow instructions from human to ask questions.\"\n",
    "exp_task_pack = TaskPackage(instruction=exp_task)\n",
    "\n",
    "manager_agent(exp_task_pack)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
